# '''
# Author: Jiahui Sun kianamustwin@163.com
# Date: 2023-09-13 20:41:24
# LastEditors: Jiahui Sun kianamustwin@163.com
# LastEditTime: 2023-09-13 20:55:11
# FilePath: /laketicv/tools/gpu_e2e/weice/cls_infer.py
# Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
# '''
# import os
# import torch
# import torch.nn.functional as F
# import torchvision.transforms as transforms
# import numpy as np
# from glob import glob
# from tqdm import tqdm
# from PIL import Image

# from mmengine.config import Config
# from mmengine.runner import load_checkpoint
# from mmpretrain.models import build_classifier

# labels = ['0', '1', '2', '3']

# def cls_infer(img_dirs):
#     config_file = "../configs/resnet18_bs256_100e_lr01_pretrained_weice_pad_4cls_v2_0803.py"
#     checkpoint_file = "../float_models/best_accuracy_top1_epoch_54.pth"
#     device = torch.device('cuda:0') 
#     cfg= Config.fromfile(config_file)
#     model = build_classifier(cfg.model)

#     checkpoint=load_checkpoint(model,checkpoint_file)
#     model = model.to(device)
#     model.eval()

#     # transform
#     transform = transforms.Compose([
#         transforms.Resize((224,224)),
#         transforms.ToTensor(),
#         transforms.Normalize(mean=[123.675/255.0, 116.28/255.0, 103.53/255.0],std=[58.395/255.0, 57.12/255.0, 57.375/255.0])
#     ])
#     result = []
#     # img_dirs = glob(os.path.join(image_path, "*.jpeg")) + glob(os.path.join(image_path, "*.jpg")) + glob(os.path.join(image_path, "*.png"))
#     for img_dir in tqdm(img_dirs):
#         with torch.no_grad():
#             image = Image.open(img_dir).convert('RGB')
#             input_tensor = transform(image).unsqueeze(0)  
#             input_tensor = input_tensor.to(device)   
#             logits = model(input_tensor).cpu()
#             pred_cls = torch.argmax(logits, dim=1).item()
#             pred_prob = F.softmax(logits[0], dim=0)
#             data = list(zip(labels,pred_prob))
#             data = [(label, value.item()) for label, value in data]
#             sorted_data = sorted(data, key=lambda x: x[1], reverse = True)
#         result.append((img_dir,sorted_data))
#     print(result)
#     return result

# if __name__ == "__main__":
#     image_path = "/nfs/algorithm/dataset/private-dataset/gary/weice/pad_cls/test/for_test/02_unknown/"
#     img_dirs = glob(os.path.join(image_path, "*.jpeg")) + glob(os.path.join(image_path, "*.jpg")) + glob(os.path.join(image_path, "*.png"))
#     cls_infer(img_dirs)

import mmcv
from mmpretrain.apis import init_model as cls_init_model
from mmpretrain.apis import inference_model as cls_inference_model
from PIL import Image
import torch
import torch.nn.functional as F
import torchvision.transforms as transforms
import numpy as np
from glob import glob
from tqdm import tqdm
import os
from mmpretrain.apis.image_classification import ImageClassificationInferencer
    

def cls_init(config_file,checkpoint_file):
    device = torch.device('cuda:0')
    config_file = config_file
    checkpoint_file = checkpoint_file
    model = cls_init_model(config_file, checkpoint_file, device='cuda:0') 
    inferencer = ImageClassificationInferencer(model)
    return inferencer

def cls_infer(inferencer,img_dirs):
    result = []
    with torch.no_grad():
        output = inferencer(img_dirs,batch_size=128)
    for i , item in enumerate(output):
        label = item['pred_label']
        score = item['pred_score']
        result.append((img_dirs[i],[label,score]))
    #print(result)
    return result